import re
import jieba
import numpy as np
from sklearn.preprocessing import LabelEncoder
from datasets import load_dataset

# 文本清洗函数
def clean_text(text):
    # 移除特殊字符
    text = re.sub(r'[^\u4e00-\u9fff\s]', '', text).strip()
    # 分词处理
    text = ' '.join(jieba.cut(text))
    return text

# 数据预处理主函数
def main():
    # 加载CrossWOZ数据集（假设已下载到本地）
    crosswoz = load_dataset(
        'json', 
        data_files={
            'train': 'data/raw/crosswoz/train.json',
            'validation': 'data/raw/crosswoz/val.json',
            'test': 'data/raw/crosswoz/test.json'
        },
        cache_dir='data/cache',  # 指定项目内缓存目录
        ignore_verifications=True  # 忽略缓存验证
    )

    # 标签编码 - 仅使用CrossWOZ数据集的标签
    intent_labels = list(set(crosswoz['train']['intent']))
    label_encoder = LabelEncoder()
    label_encoder.fit(intent_labels)
    np.save('data/label_classes.npy', label_encoder.classes_)

    # 数据预处理函数
    def preprocess_function(examples):
        # 文本清洗
        examples['text'] = [clean_text(text) for text in examples['text']]
        # 标签编码
        examples['label'] = label_encoder.transform(examples['intent'])
        return examples

    # 应用预处理
    processed_crosswoz = crosswoz.map(preprocess_function, batched=True)

    # 保存处理后的数据 - 仅使用CrossWOZ数据集
    processed_crosswoz['train'].to_json('data/processed/train.json')
    processed_crosswoz['validation'].to_json('data/processed/val.json')
    processed_crosswoz['test'].to_json('data/processed/test.json')

if __name__ == "__main__":
    main()